import os
import math
import argparse
import csv
import random
from typing import List, Tuple

import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader


FEATURE_COLS = ["s_det", "sim1", "margin", "area_ratio", "aspect_ratio", "best_iou", "delta_sim1"]
LABEL_COL = "y"


class LWGDataset(Dataset):
    def __init__(self, rows: List[Tuple[List[float], float]]):
        self.rows = rows

    def __len__(self):
        return len(self.rows)

    def __getitem__(self, idx):
        x, y = self.rows[idx]
        return torch.tensor(x, dtype=torch.float32), torch.tensor([y], dtype=torch.float32)


class GateMLP(nn.Module):
    def __init__(self, in_dim: int = 7, hidden: List[int] = [64, 16]):
        super().__init__()
        layers: List[nn.Module] = []
        last = in_dim
        for h in hidden:
            layers += [nn.Linear(last, h), nn.ReLU(inplace=True)]
            last = h
        layers += [nn.Linear(last, 1), nn.Sigmoid()]
        self.net = nn.Sequential(*layers)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x).view(-1, 1)


def set_seed(seed: int):
    random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)


def load_rows(csv_path: str) -> List[Tuple[List[float], float]]:
    rows: List[Tuple[List[float], float]] = []
    with open(csv_path) as f:
        reader = csv.DictReader(f)
        for row in reader:
            try:
                best_iou = float(row.get("best_iou", "0") or 0.0)
                if not math.isfinite(best_iou):
                    best_iou = 0.0
                best_iou = max(0.0, min(1.0, best_iou))
                delta_sim = float(row.get("delta_sim1", "0") or 0.0)
                if not math.isfinite(delta_sim):
                    delta_sim = 0.0
                delta_sim = max(-2.0, min(2.0, delta_sim))
                x = [
                    float(row["s_det"]),
                    float(row["sim1"]) if row["sim1"] != "" else 0.0,
                    float(row["margin"]) if row["margin"] != "" else 0.0,
                    max(0.0, min(1.0, float(row["area_ratio"]))),
                    float(row["aspect_ratio"]) if row["aspect_ratio"] != "" else 1.0,
                    best_iou,
                    delta_sim,
                ]
                y = float(row[LABEL_COL])
                rows.append((x, y))
            except Exception:
                continue
    return rows


def train_one(
    csv_path: str,
    out_path: str,
    epochs: int = 5,
    batch_size: int = 8192,
    lr: float = 1e-3,
    device: str = "cuda",
    val_ratio: float = 0.1,
    patience: int = 5,
    seed: int = 42,
):
    set_seed(seed)
    rows = load_rows(csv_path)
    assert len(rows) > 0, f"No data loaded from {csv_path}"

    val_size = int(len(rows) * max(0.0, min(val_ratio, 0.5)))
    if val_size > 0 and len(rows) - val_size < 32:
        val_size = 0
    random = torch.Generator().manual_seed(seed)
    dataset = LWGDataset(rows)
    if val_size > 0:
        train_size = len(rows) - val_size
        train_ds, val_ds = torch.utils.data.random_split(dataset, [train_size, val_size], generator=random)
    else:
        train_ds = dataset
        val_ds = None

    train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True, drop_last=False)
    val_loader = DataLoader(val_ds, batch_size=batch_size, shuffle=False, drop_last=False) if val_ds else None

    model = GateMLP(in_dim=len(FEATURE_COLS), hidden=[64, 16]).to(device)
    opt = torch.optim.AdamW(model.parameters(), lr=lr)
    loss_fn = nn.BCELoss()

    best_val = float("inf")
    best_state = None
    wait = 0

    for epoch in range(epochs):
        model.train()
        total_loss = 0.0
        total_samples = 0
        for x, y in train_loader:
            x = x.to(device)
            y = y.to(device)
            pred = model(x)
            loss = loss_fn(pred, y)
            opt.zero_grad()
            loss.backward()
            opt.step()
            total_loss += float(loss.item()) * x.shape[0]
            total_samples += x.shape[0]
        train_loss = total_loss / max(1, total_samples)

        if val_loader is not None:
            model.eval()
            val_loss = 0.0
            val_samples = 0
            with torch.no_grad():
                for x, y in val_loader:
                    x = x.to(device)
                    y = y.to(device)
                    pred = model(x)
                    loss = loss_fn(pred, y)
                    val_loss += float(loss.item()) * x.shape[0]
                    val_samples += x.shape[0]
            val_loss /= max(1, val_samples)
            print(f"[epoch {epoch+1}/{epochs}] train={train_loss:.6f} val={val_loss:.6f} N_train={total_samples} N_val={val_samples}")

            if val_loss < best_val - 1e-6:
                best_val = val_loss
                wait = 0
                best_state = {k: v.cpu() for k, v in model.state_dict().items()}
            else:
                wait += 1
                if wait >= max(1, patience):
                    print(f"[early stop] patience reached at epoch {epoch+1}")
                    break
        else:
            print(f"[epoch {epoch+1}/{epochs}] train={train_loss:.6f} N_train={total_samples}")

    if best_state is not None:
        model.load_state_dict(best_state)

    payload = {
        "in_dim": len(FEATURE_COLS),
        "hidden": [64, 16],
        "state_dict": {('net.'+k if not k.startswith('net.') else k): v.cpu() for k, v in model.state_dict().items()},
    }
    os.makedirs(os.path.dirname(out_path), exist_ok=True)
    torch.save(payload, out_path)
    print(f"[OK] Saved LWG to {out_path}")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--csv", type=str, required=True, help="Path to train_lwg.csv")
    ap.add_argument("--out", type=str, default="/home/yves/MOTIP/pretrains/lwg_dancetrack.pt")
    ap.add_argument("--epochs", type=int, default=5)
    ap.add_argument("--batch-size", type=int, default=8192)
    ap.add_argument("--lr", type=float, default=1e-3)
    ap.add_argument("--device", type=str, default="cuda")
    ap.add_argument("--val-ratio", type=float, default=0.1)
    ap.add_argument("--patience", type=int, default=5)
    ap.add_argument("--seed", type=int, default=42)
    args = ap.parse_args()
    train_one(
        csv_path=args.csv,
        out_path=args.out,
        epochs=args.epochs,
        batch_size=args.batch_size,
        lr=args.lr,
        device=args.device,
        val_ratio=args.val_ratio,
        patience=args.patience,
        seed=args.seed,
    )


if __name__ == "__main__":
    main()
